from rest_framework import viewsets, permissions, status
from rest_framework.permissions import BasePermission
from rest_framework.decorators import action
from rest_framework.response import Response
import json

from .models import EnglishText
from .serializers import EnglishTextSerializer

from transformer.inference import inference_model


class AllowWriteResource(BasePermission):
    def has_permission(self, request, view):
        return request.user.groups.filter(name='管理员').exists() or request.user.groups.filter(name='研究员').exists()


class EnglishTextViewSet(viewsets.ModelViewSet):
    queryset = EnglishText.objects.all()
    serializer_class = EnglishTextSerializer

    def get_permissions(self):
        if self.action in ['create', 'destroy', 'update']:
            permission_classes = [permissions.IsAuthenticated, AllowWriteResource]
        else:
            permission_classes = [permissions.IsAuthenticated]
        return [permission() for permission in permission_classes]

    @action(detail=False, methods=['post'])
    def translate(self, request):
        text_en = request.data.get('text_en')
        if not text_en:
            return Response({'error': 'text_en is required'}, status=status.HTTP_400_BAD_REQUEST)

        src_tokens, src_words, output_words, attentions = inference_model(text_en)

        return Response({'src_tokens': src_tokens, 'src_words': src_words, 'output_words': output_words, 'attentions':  json.dumps(attentions.numpy().tolist())})
